import os
import pandas as pd
from pathlib import Path

"""
删除mosei—text中过多的部分 以加快数据处理
"""
def process_features_and_filter_csv(feature_folder, target_csv):
    """
    处理特征文件并过滤CSV文件

    参数：
    feature_folder : str - 包含特征文件的文件夹路径
    target_csv : str - 需要过滤的目标CSV文件路径
    """

    # 第一步：处理特征文件名
    feature_files = []
    valid_names = set()

    try:
        # 遍历特征文件夹
        for filename in os.listdir(feature_folder):
            # 检查是否以_features.csv结尾
            if filename.endswith("_features.csv"):
                # 提取基础名称并添加到集合
                base_name = filename.rsplit("_features.csv", 1)[0]
                valid_names.add(base_name)
                feature_files.append(filename)

        print(f"找到 {len(valid_names)} 个有效特征文件")

    except FileNotFoundError:
        print(f"错误：特征文件夹 {feature_folder} 不存在")
        return
    except Exception as e:
        print(f"读取特征文件夹时发生错误: {str(e)}")
        return

    # 第二步：处理目标CSV文件
    try:
        # 读取CSV文件
        df = pd.read_csv(target_csv)

        # 检查是否存在name列
        if "name" not in df.columns:
            print("错误：CSV文件中缺少 'name' 列")
            return

        # 过滤数据
        original_count = len(df)
        filtered_df = df[df["name"].isin(valid_names)]
        new_count = len(filtered_df)

        print(f"原始数据量: {original_count}")
        print(f"过滤后数据量: {new_count}")
        print(f"移除数据量: {original_count - new_count}")

        # 备份原文件
        backup_path = Path(target_csv).with_suffix(".bak.csv")
        df.to_csv(backup_path, index=False)
        print(f"已创建备份文件: {backup_path}")

        # 保存过滤后的文件
        filtered_df.to_csv(target_csv, index=False)
        print(f"已更新目标文件: {target_csv}")

    except FileNotFoundError:
        print(f"错误：目标CSV文件 {target_csv} 不存在")
    except pd.errors.EmptyDataError:
        print("错误：CSV文件内容为空或格式不正确")
    except Exception as e:
        print(f"处理CSV文件时发生错误: {str(e)}")


if __name__ == "__main__":
    # 配置路径（请根据实际情况修改）
    FEATURE_FOLDER = "../../data/out_feature_mosei"
    TARGET_CSV = "../../data/cmu_mosei/merged_data_clean_max.csv"

    # 执行处理
    process_features_and_filter_csv(FEATURE_FOLDER, TARGET_CSV)